…transitions

Adds a tutorial section demonstrating a genuinely continuous-time SDE
(ContinuousTimeStateEvolution + Discretizer(discretize=euler_maruyama))
composed with DiscreteControlLoopSimulator, showing the loop only ever sees
discrete times while the underlying dynamics are a true SDE solved between
them.

Building this surfaced a real bug: the bootstrap FilterUpdate call left
t_prev defaulting to t (dt=0). For fixed-covariance transitions this is
harmless, but for a dt-scaled transition (like the Euler-Maruyama
discretization here), it constructs a zero-covariance distribution whose
NaN log-density leaks through the gradient in EKF's Taylor linearization
(jnp.where evaluates both branches, unlike jax.lax.cond) -- corrupting the
filtered state and, downstream, the policy's control and the simulated
trajectory itself. Fixed by giving the bootstrap call a non-degenerate
t_prev, borrowing the width of the first real interval (matching
compute_cuthbert_filter's own dummy-row convention).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…ccuracy

Documents how the continuous-time integration is actually defined: Discretizer
takes exactly one Euler-Maruyama step over the whole gap between requested
times (needed so filtering gets an explicit one-step transition density),
unlike SDESimulator/solve_sde's substepped solvers, which only support pure
forward simulation without filtering.

Adds a section 6 demonstrating the same recipe on a genuinely nonlinear 2D
system with drift = A x^2 + u: x=0 is an unstable equilibrium (x^2 >= 0 always
pushes away from the origin), so the uncontrolled run diverges while the same
linear feedback policy as before stabilizes it locally. Uses a finer time grid
(dt=0.1 vs 1.0) since the one-step EM approximation needs a small enough gap
to stay accurate for a nonlinear drift.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…rt_filter_update

35 tests across 7 groups: filter_state_mean unit tests, compute_cuthbert_filter_update
correctness (including regression tests for the dt=0 EKF NaN bug), validation/error
paths, end-to-end shape/output-key tests (including the stateless-policy regression),
behavioral/control correctness (closed-loop stabilization, the control-index
convention, determinism, eqx.Module policies), continuous-time/Discretizer
composition, and black-box transition compatibility.

Surfaced one more real gap while writing these: a genuinely nested-pytree
policy_state_init (e.g. a dict of arrays) can't actually round-trip today --
BaseSimulator's shared _run_single_member_simulation enforces a
dict[str, Array] | None return type via jaxtyping at runtime, so
policy_states must stay a flat Array. Documented in the corresponding test
rather than widening the shared base class's type contract, which is out of
scope here.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…ulator there

Consolidates control-loop + control-policy code under dynestyx/control/:
moves discrete_controller_simulators.py there, and adds a basic MPPI
(Model Predictive Path Integral) controller (dynestyx/control/mppi.py) that
plugs into the same control_policy= slot. Deliberately simple: samples
candidate control sequences as Gaussian perturbations around a nominal
sequence, rolls each through a user-supplied dynamics_model, and returns
the softmax-weighted mean control (standard MPPI weighting). dynamics_model
can be batched (one call handles all samples) or not (driven via
jax.lax.map, so it never needs to support batching itself).

MPPI needs fresh randomness every step, which PolicyCallable's signature
didn't provide -- added a key parameter to __call__(x_hat, s, key), passed
in by DiscreteControlLoopSimulator's per-step scan. This also keeps MPPI's
own policy state a flat array (just the nominal control sequence), avoiding
a real limitation found while testing last time: a nested-pytree
policy_state_init can't actually be recorded as output today, since
BaseSimulator's shared return-type check requires flat Array values.

Verified: existing linear-feedback policies (tests, tutorial notebook)
updated for the new signature, full regression suite still passes
(117 tests), and a smoke test confirms MPPI stabilizes a marginally-unstable
linear system identically under both batched and non-batched dynamics_model.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

# Conflicts:
#	dynestyx/inference/integrations/cuthbert/discrete_filter.py

…lation)

main/control_v2 landed a major refactor (commit 3fb1cd7, "Numpyro-Free
Usage") that deleted dynestyx/simulators.py entirely, replacing it with a
new dynestyx/simulation/ package: BaseSimulator subclasses now implement a
public `simulate(dynamics, *, rng_key, ...) -> SimulatedResult` (not
`_simulate(name, ..., obs_times=..., ...) -> dict`), take rng_key directly
instead of calling numpyro.prng_key() internally, and numpyro site
registration happens via a generic deferred callback (dataclasses.fields()
iteration over the returned SimulatedResult) rather than base-class magic
over a free-form dict.

Ported DiscreteControlLoopSimulator to the new contract:
- Plain class with explicit __init__ (matching DiscreteTimeSimulator's new
  shape) instead of a bare @dataclasses.dataclass.
- New ControlledSimulatedResult(SimulatedResult) subclass carrying the
  extra controls/filtered_states_mean/policy_states fields -- confirmed by
  reading base.py directly that _run_single_member_simulation's common path
  is a bare `return self.simulate(...)` with no reconstruction, so subclass
  fields flow through untouched and get registered generically the same way
  as SimulatedResult's own fields (None-valued fields are auto-skipped).
- Dropped the now-redundant manual `numpyro.prng_key()`/obs_values checks:
  the base class already validates/rejects these before simulate() is ever
  called.

compute_cuthbert_filter_update/_build_cuthbert_filter_obj (added to
discrete_filter.py on the `control` branch) merged in via `git merge control`
with a single trivial conflict (re-adding the `jax.random` import) --
confirmed via git merge-tree dry run and direct diffing that every symbol
they depend on exists byte-identical post-refactor, just re-imported from
dynestyx.inference.configs.filter instead of dynestyx.inference.filter_configs.
Discretizer required no changes at all (already ported upstream).

Verified: full regression suite passes on control_v2 (118 tests: this
file's 35 plus test_filters.py/test_filter_simulator.py/test_discretizers.py/
test_predictive_filter_simulator_shapes.py), the tutorial notebook
re-executes cleanly end-to-end, and MPPI produces numerically identical
results to its control-branch run under both batched and non-batched
dynamics_model.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Section 7 demonstrates dynestyx.control.MPPI plugged into the same
control_policy= slot as the earlier linear-feedback policy, on the same 1D
system from section 1. Builds a planning rollout directly from
dynamics.state_evolution's deterministic mean (a standard MPPI
simplification -- the planner doesn't need a faithful stochastic
simulation, just a reasonable prediction of where a candidate control
sequence leads) and a simple quadratic loss, then compares against the
K=0 no-control baseline already computed in section 3.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

DynamicalModel.state_evolution is declared as the broader
ContinuousTimeStateEvolution | DiscreteStateTransition union, so `ty`
couldn't statically narrow continuous_dynamics.state_evolution to
StochasticContinuousTimeStateEvolution before passing it to
euler_maruyama() (which requires the narrower type) -- even though this is
exactly how Discretizer._sample_ds itself resolves it at runtime. Added the
same isinstance narrowing check used there.

Also includes an incidental ruff-format tweak to mppi.py (collapsing a
one-line call that had been wrapped unnecessarily).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…into control_v2

Section 3 (differentiating through the closed loop) computed grad_K via
jax.grad but never independently verified it. Added a cell checking a
central finite difference on K[0, 0] against the autodiff value -- confirms
the gradient through the whole closed loop (policy -> transition ->
observation -> filter update) is actually correct, not just plausible.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Section 3's rollout_final_state_norm now calls sim.simulate(dynamics,
rng_key=..., predict_times=...) directly instead of going through
dsx.sample/numpyro.handlers.seed -- unneeded here since there's no
conditioning/MCMC, and it makes the PRNG key an explicit argument instead
of an implicit fixed rng_seed=0.

More importantly, this fixes a real issue in section 4's training loop:
every epoch previously reused the same fixed seed, so gradient descent could
overfit K to one specific noise realization rather than the underlying
dynamics. Now each epoch draws a fresh key via jax.random.split, so K_opt
has to generalize across realizations. Verified the loop still converges to
a sensible stabilizing gain with this fresh-key-per-epoch setup before
editing the notebook. Section 4's final comparison plot (run()/dsx.sample)
is left unchanged -- there, reusing the same seed for both the unoptimized
and optimized rollouts is what makes the before/after comparison fair.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

run() (updated separately) now returns a ControlledSimulatedResult from
sim.simulate() directly instead of a numpyro trace dict, so the plotting
cell needed the equivalent attribute-access update: trace["loop_X"]["value"]
-> result.X. Also removed a leftover duplicate plotting cell still using the
old trace-dict access, which was left after the updated one and caused the
notebook to fail on execution.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Replaces the with sim: dsx.sample(...) + numpyro.handlers.seed/trace pattern
across every section (discrete demo, SDE, nonlinear 2D, MPPI) with direct
sim.simulate(dynamics, rng_key=key, predict_times=...) calls returning
ControlledSimulatedResult objects -- no NumPyro handler or model function
needed, since every section here is a pure generative rollout with nothing
to condition on. Renamed trace_* variables to result_* throughout, since
they're no longer NumPyro trace dicts, and updated all downstream plotting
cells from trace["name_field"]["value"] to result.field attribute access.

For the continuous-time sections, this also removes the Discretizer handler
composition: instead of wrapping dsx.sample in `with Discretizer(...)`, the
continuous-time DynamicalModel is discretized once up front via
euler_maruyama(state_evolution) into a plain discrete-time DynamicalModel,
which is then simulated exactly like any other discrete-time model. Updated
the surrounding markdown to match (no more handler-chain explanation needed).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Unified non-linear dynamics, black box and partial observations example

Rework MPPI to take one-step dynamics; complete and stabilize mpc_demo

- MPPI now takes dynamics: DynamicalModel (one-step transition) instead of
  a hand-rolled full-rollout function; it builds the horizon-unrolling and
  n_samples batching internally via jax.lax.scan/vmap, using .mean when
  available and falling back to .sample() (own internal key) for black-box
  transitions.
- All fields but dynamics/loss_fn now have defaults (horizon=10,
  noise_std=1.0, n_samples=20, dt=1.0, temperature=1.0, batched=True,
  seed=0).
- stability: Non-finite (nan/inf) rollout losses are clamped before the softmax
  weighting, fixing NaN control
- mpc_demo.ipynb: completed (previously referenced undefined names, never
  ran), updated to the new MPPI API, added MPC/MPPI citations.